import numpy as np
import os
import time


def generate_random_tsp_data(num_cities: int, filename_prefix: str = "random_tsp"):
    """
    生成随机TSP数据集并保存到文件，文件名包含节点数量和时间戳

    Args:
        num_cities: 城市数量
        filename_prefix: 文件名前缀（默认为"random_tsp"）
    """
    # 生成随机城市坐标（x, y）
    locations = np.random.rand(num_cities, 2) * 100  # 坐标范围在[0, 100]

    output_dir = "../data/random"
    os.makedirs(output_dir, exist_ok=True)

    timestamp = int(time.time())
    filename = f"{filename_prefix}_{num_cities}_{timestamp}.tsp"
    filepath = os.path.join(output_dir, filename)

    with open(filepath, 'w') as f:
        f.write("NAME : random_tsp\n")
        f.write("COMMENT : Randomly generated TSP instance\n")
        f.write(f"TYPE : TSP\n")
        f.write(f"DIMENSION : {num_cities}\n")
        f.write("EDGE_WEIGHT_TYPE : EUC_2D\n")
        f.write("\n")
        f.write("NODE_COORD_SECTION\n")

        for i in range(num_cities):
            f.write(f"{i + 1} {locations[i][0]:.6f} {locations[i][1]:.6f}\n")

        f.write("EOF\n")

    print(f"随机TSP数据集已生成并保存至: {filepath}")
    return locations


if __name__ == "__main__":
    # 生成5个城市的TSP数据
    cities = generate_random_tsp_data(10)
